# frozen_string_literal: true

require 'spec_helper'

RSpec.describe Gitlab::BackgroundMigration::BackfillVulnerabilityFindingRiskScores,
  feature_category: :vulnerability_management do
  let(:vulnerability_finding_risk_scores) { table(:vulnerability_finding_risk_scores, database: :sec) }
  let(:findings) { table(:vulnerability_occurrences, database: :sec) }
  let(:identifiers) { table(:vulnerability_identifiers, database: :sec) }
  let(:finding_identifiers) { table(:vulnerability_occurrence_identifiers, database: :sec) }
  let(:scanners) { table(:vulnerability_scanners, database: :sec) }
  let(:cve_enrichments) { table(:pm_cve_enrichment) }
  let(:organizations) { table(:organizations) }
  let(:namespaces) { table(:namespaces) }
  let(:projects) { table(:projects) }

  # Severity enum values from Enums::Vulnerability::SEVERITY_LEVELS
  let(:severity_info) { 1 }
  let(:severity_unknown) { 2 }
  let(:severity_low) { 4 }
  let(:severity_medium) { 5 }
  let(:severity_high) { 6 }
  let(:severity_critical) { 7 }

  let!(:organization) { organizations.create!(name: 'test-org', path: 'test-org') }
  let!(:group_namespace) do
    namespaces.create!(name: 'test-group', path: 'test-group', organization_id: organization.id)
  end

  let!(:project_namespace) do
    namespaces.create!(name: 'test-project', path: 'test-project', organization_id: organization.id)
  end

  let!(:project) do
    projects.create!(
      name: 'test-project',
      path: 'test-project',
      namespace_id: group_namespace.id,
      project_namespace_id: project_namespace.id,
      organization_id: organization.id
    )
  end

  let!(:scanner) do
    scanners.create!(
      project_id: project.id,
      external_id: 'test-scanner',
      name: 'Test Scanner'
    )
  end

  def args
    {
      start_id: findings.minimum(:id),
      end_id: findings.maximum(:id),
      batch_table: :vulnerability_occurrences,
      batch_column: :id,
      sub_batch_size: 100,
      pause_ms: 0,
      connection: SecApplicationRecord.connection
    }
  end

  subject(:perform_migration) { described_class.new(**args).perform }

  def create_finding(severity:, cve: nil, epss_score: nil, is_known_exploit: false, with_risk_score: false)
    identifier = identifiers.create!(
      project_id: project.id,
      external_id: cve || "CVE-#{SecureRandom.hex(4)}",
      external_type: cve ? 'CVE' : 'OTHER',
      name: cve || "Identifier-#{SecureRandom.hex(4)}",
      fingerprint: SecureRandom.hex(20)
    )

    finding = findings.create!(
      project_id: project.id,
      scanner_id: scanner.id,
      primary_identifier_id: identifier.id,
      location_fingerprint: SecureRandom.hex(20),
      uuid: SecureRandom.uuid,
      name: 'Test Finding',
      severity: severity,
      report_type: 1,
      raw_metadata: '{}',
      metadata_version: 'test:1.0'
    )

    finding_identifiers.create!(
      occurrence_id: finding.id,
      identifier_id: identifier.id
    )

    if cve && (epss_score || is_known_exploit)
      cve_enrichments.create!(
        cve: cve,
        epss_score: epss_score,
        is_known_exploit: is_known_exploit
      )
    end

    if with_risk_score
      vulnerability_finding_risk_scores.create!(
        finding_id: finding.id,
        project_id: project.id,
        risk_score: 0.0
      )
    end

    finding
  end

  describe '#perform' do
    context 'with critical severity and no CVE enrichment' do
      let!(:finding) { create_finding(severity: severity_critical) }

      it 'creates risk score record with 0.6 (base score for critical)' do
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        risk_score = vulnerability_finding_risk_scores.find_by(finding_id: finding.id)
        expect(risk_score.risk_score).to eq(0.6)
        expect(risk_score.project_id).to eq(project.id)
      end
    end

    context 'with high severity and no CVE enrichment' do
      let!(:finding) { create_finding(severity: severity_high) }

      it 'creates risk score record with 0.4 (base score for high)' do
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(0.4)
      end
    end

    context 'with medium severity and no CVE enrichment' do
      let!(:finding) { create_finding(severity: severity_medium) }

      it 'creates risk score record with 0.2 (base score for medium)' do
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(0.2)
      end
    end

    context 'with low severity and no CVE enrichment' do
      let!(:finding) { create_finding(severity: severity_low) }

      it 'creates risk score record with 0.05 (base score for low)' do
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(0.05)
      end
    end

    context 'with CVE enrichment having high EPSS score' do
      let!(:finding) do
        create_finding(
          severity: severity_high,
          cve: 'CVE-2023-12345',
          epss_score: 0.8,
          is_known_exploit: false
        )
      end

      it 'creates risk score with EPSS modifier' do
        # base_score (0.4) + epss_base_modifier (0.8 * 0.3 = 0.24) + epss_bonus (0.2) = 0.84
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to be_within(0.01).of(0.84)
      end
    end

    context 'with CVE enrichment having known exploit' do
      let!(:finding) do
        create_finding(
          severity: severity_high,
          cve: 'CVE-2023-54321',
          epss_score: 0.0,
          is_known_exploit: true
        )
      end

      it 'creates risk score with KEV modifier' do
        # base_score (0.4) + kev_modifier (0.3) = 0.7
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(0.7)
      end
    end

    context 'with CVE enrichment having both high EPSS and known exploit' do
      let!(:finding) do
        create_finding(
          severity: severity_critical,
          cve: 'CVE-2023-99999',
          epss_score: 0.9,
          is_known_exploit: true
        )
      end

      it 'creates risk score capped at 1.0' do
        # base_score (0.6) + epss_base_modifier (0.9 * 0.3 = 0.27) + epss_bonus (0.2) + kev_modifier (0.3) = 1.37
        # Capped at 1.0
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(1.0)
      end
    end

    context 'with multiple findings' do
      let!(:finding1) { create_finding(severity: severity_critical) }
      let!(:finding2) { create_finding(severity: severity_high) }
      let!(:finding3) { create_finding(severity: severity_medium) }

      it 'creates all risk scores correctly' do
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(3)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding1.id).risk_score).to eq(0.6)
        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding2.id).risk_score).to eq(0.4)
        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding3.id).risk_score).to eq(0.2)
      end
    end

    context 'when risk score record already exists' do
      let!(:finding) { create_finding(severity: severity_critical, with_risk_score: true) }

      it 'updates the existing risk score record' do
        expect { perform_migration }
          .not_to change { vulnerability_finding_risk_scores.count }

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(0.6)
      end
    end

    context 'when finding has updated severity' do
      let!(:finding) { create_finding(severity: severity_low, with_risk_score: true) }

      before do
        findings.find(finding.id).update!(severity: severity_critical)
      end

      it 'updates risk score to reflect new severity' do
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score }
          .from(0.0).to(0.6)
      end
    end

    context 'when CVE enrichment has nil epss_score' do
      let!(:finding) do
        create_finding(
          severity: severity_high,
          cve: 'CVE-2023-00000',
          epss_score: nil,
          is_known_exploit: false
        )
      end

      it 'treats nil epss_score as 0 and calculates base score only' do
        # base_score (0.4) + epss_base_modifier (0) + epss_bonus (0) = 0.4
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(0.4)
      end
    end

    context 'when finding has no CVE identifiers' do
      let!(:finding) { create_finding(severity: severity_high) }

      it 'creates risk score with base score only' do
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(0.4)
      end
    end

    context 'with unknown severity' do
      let!(:finding) { create_finding(severity: severity_unknown) }

      it 'creates risk score with 0.2 (base score for unknown)' do
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(0.2)
      end
    end

    context 'with info severity' do
      let!(:finding) { create_finding(severity: severity_info) }

      it 'creates risk score with 0 (base score for info)' do
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(0.0)
      end
    end

    context 'for N+1 queries' do
      def run_migration
        described_class.new(**args).perform
      end

      it 'does not create N+1 queries when loading CVE enrichments' do
        create_finding(
          severity: severity_high,
          cve: "CVE-2023-1001",
          epss_score: 0.5,
          is_known_exploit: false
        )

        control = ActiveRecord::QueryRecorder.new(skip_cached: false) do
          run_migration
        end

        5.times do |i|
          create_finding(
            severity: severity_high,
            cve: "CVE-2023-#{2000 + i}",
            epss_score: 0.6,
            is_known_exploit: true
          )
        end

        expect { run_migration }.not_to exceed_query_limit(control)
      end
    end

    context 'with EPSS score edge cases' do
      it 'applies correct bonus for EPSS score at 0.1 threshold' do
        finding = create_finding(
          severity: severity_high,
          cve: 'CVE-2023-10000',
          epss_score: 0.1,
          is_known_exploit: false
        )

        # base_score (0.4) + epss_base_modifier (0.1 * 0.3 = 0.03) + epss_bonus (0.1) = 0.53
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(0.53)
      end

      it 'applies correct bonus for EPSS score at 0.5 threshold' do
        finding = create_finding(
          severity: severity_medium,
          cve: 'CVE-2023-50000',
          epss_score: 0.5,
          is_known_exploit: false
        )

        # base_score (0.2) + epss_base_modifier (0.5 * 0.3 = 0.15) + epss_bonus (0.2) = 0.55
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(0.55)
      end

      it 'applies no bonus for EPSS score below 0.1' do
        finding = create_finding(
          severity: severity_high,
          cve: 'CVE-2023-09999',
          epss_score: 0.09,
          is_known_exploit: false
        )

        # base_score (0.4) + epss_base_modifier (0.09 * 0.3 = 0.027) + epss_bonus (0) = 0.427
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to be_within(0.001).of(
          0.427)
      end
    end

    context 'with multiple CVE identifiers' do
      it 'uses the first CVE enrichment found' do
        identifier1 = identifiers.create!(
          project_id: project.id,
          external_id: 'CVE-2023-11111',
          external_type: 'CVE',
          name: 'CVE-2023-11111',
          fingerprint: SecureRandom.hex(20)
        )

        identifier2 = identifiers.create!(
          project_id: project.id,
          external_id: 'CVE-2023-22222',
          external_type: 'CVE',
          name: 'CVE-2023-22222',
          fingerprint: SecureRandom.hex(20)
        )

        finding = findings.create!(
          project_id: project.id,
          scanner_id: scanner.id,
          primary_identifier_id: identifier1.id,
          location_fingerprint: SecureRandom.hex(20),
          uuid: SecureRandom.uuid,
          name: 'Test Finding',
          severity: severity_high,
          report_type: 1,
          raw_metadata: '{}',
          metadata_version: 'test:1.0'
        )

        finding_identifiers.create!(
          occurrence_id: finding.id,
          identifier_id: identifier1.id
        )

        finding_identifiers.create!(
          occurrence_id: finding.id,
          identifier_id: identifier2.id
        )

        # Only create enrichment for the second CVE
        cve_enrichments.create!(
          cve: 'CVE-2023-22222',
          epss_score: 0.8,
          is_known_exploit: true
        )

        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        # Should use the enrichment data from CVE-2023-22222
        # base_score (0.4) + epss_base_modifier (0.8 * 0.3 = 0.24) + epss_bonus (0.2) + kev_modifier (0.3) =
        # 1.14 capped at 1.0
        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(1.0)
      end
    end

    context 'with combined enrichment modifiers' do
      it 'correctly calculates score with medium EPSS and known exploit' do
        finding = create_finding(
          severity: severity_low,
          cve: 'CVE-2023-33333',
          epss_score: 0.3,
          is_known_exploit: true
        )

        # base_score (0.05) + epss_base_modifier (0.3 * 0.3 = 0.09) + epss_bonus (0.1) + kev_modifier (0.3) = 0.54
        expect { perform_migration }
          .to change { vulnerability_finding_risk_scores.count }.by(1)

        expect(vulnerability_finding_risk_scores.find_by(finding_id: finding.id).risk_score).to eq(0.54)
      end
    end

    context 'when sub-batch is empty' do
      it 'handles empty batches gracefully' do
        finding = create_finding(severity: severity_high)
        findings.where(id: finding.id).delete_all

        expect { perform_migration }.not_to raise_error
        expect(vulnerability_finding_risk_scores.count).to eq(0)
      end
    end
  end
end
